"""
水库群供水调度优化模型数据加载器
支持通用化的JSON配置文件和数据文件
"""

import json
import os
import argparse
from datetime import datetime
from typing import Dict, List, Any, Optional
from loguru import logger
from optmodel import log_config  # 导入日志配置

def load_json(filename: str) -> Dict[str, Any]:
    """
    加载JSON文件
    
    Args:
        filename: JSON文件路径
        
    Returns:
        JSON数据字典
    """
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            data = json.load(f)
        logger.info(f"成功加载JSON文件: {filename}")
        return data
    except FileNotFoundError:
        logger.error(f"文件未找到: {filename}")
        raise FileNotFoundError(f"文件未找到: {filename}")
    except json.JSONDecodeError as e:
        logger.error(f"JSON格式错误: {e}")
        raise ValueError(f"JSON格式错误: {e}")

def extract_config_data(config: Dict[str, Any]) -> Dict[str, Any]:
    """
    提取配置数据中的关键信息
    
    Args:
        config: 配置数据字典
        
    Returns:
        提取的配置数据
    """
    return {
        'general_parameters': config.get('general_parameters', {}),
        'input_folder': config.get('general_parameters', {}).get('input_folder', 'input'),
        'output_folder': config.get('general_parameters', {}).get('output_folder', 'output'),
        'output_filename': config.get('general_parameters', {}).get('output_filename', 'results')
    }

def parse_cli_arguments(params_list: List[str]) -> argparse.Namespace:
    """
    解析命令行参数
    
    Args:
        params_list: 参数名称列表
        
    Returns:
        解析后的命令行参数
    """
    parser = argparse.ArgumentParser(
        description='水库群供水调度优化模型命令行参数解析器'
    )
    for param in params_list:
        parser.add_argument(
            f'--{param}', type=str, default=None,
            help=f'输入参数后缀: {param}'
        )
    return parser.parse_args()

def initialize_environment(config_files: Dict[str, str]) -> Dict[str, Any]:
    """
    初始化环境，加载配置数据，设置日志，处理输入参数
    
    Args:
        config_files: 包含配置文件路径的字典
        
    Returns:
        包含全局参数的字典
    """
    config_filename = config_files['config_filename']
    filepath = config_files.get('filepath', os.getcwd())
    
    # 加载配置数据
    config_data = load_json(config_filename)
    required_config_data = extract_config_data(config_data)
    
    # 确定输入文件夹路径
    input_folder = str(config_data['general_parameters']['input_folder'])
    input_filepath = os.path.join(filepath, input_folder)
    
    # 确定输出文件夹路径
    output_folder = os.path.join(filepath, str(config_data['general_parameters']['output_folder']))
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        logger.info(f"创建输出文件夹: {output_folder}")
    
    # 确定输出文件名
    output_filename = os.path.join(output_folder, str(config_data['general_parameters']['output_filename']))
    
    params = {
        'config': config_data,
        'input_filepath': input_filepath,
        'output_filename': output_filename,
        'output_folder': output_folder
    }
    params.update(required_config_data)
    
    return params

def get_output_filename_from_config(config_path: str = "config.json") -> str:
    """
    从配置文件获取输出文件名
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        输出文件名
    """
    # 设置项目根目录路径
    if not os.path.isabs(config_path):
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_path = os.path.join(project_root, config_path)
    
    config_files = {
        'config_filename': config_path,
        'filepath': os.path.dirname(config_path)
    }
    
    params = initialize_environment(config_files)
    return params['output_filename']

def load_all_json_data(config_file: str = "config.json") -> Dict[str, Any]:
    """
    加载所有JSON数据文件（config, runoff, demand, obj_cons）
    
    Args:
        config_file: 配置文件路径
        
    Returns:
        包含所有原始数据的字典
    """
    # 确定项目根目录
    if not os.path.isabs(config_file):
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_file = os.path.join(project_root, config_file)
    else:
        project_root = os.path.dirname(config_file)
    
    logger.info(f"从配置文件加载数据: {config_file}")
    
    # 加载主配置文件
    config = load_json(config_file)
    
    # 从配置文件中获取数据文件路径
    forecast_files = config.get('forecast_data_files', {})
    input_folder = config.get('general_parameters', {}).get('input_folder', 'input')
    
    # 辅助函数：处理文件路径
    def get_full_path(filename: str) -> str:
        if os.path.isabs(filename):
            return filename
        elif not os.path.dirname(filename):
            # 只是文件名，从input文件夹查找
            return os.path.join(project_root, input_folder, filename)
        else:
            # 相对路径
            return os.path.join(project_root, filename)
    
    # 加载径流数据
    runoff_file = forecast_files.get('runoff_data', 'runoff.json')
    runoff_data = load_json(get_full_path(runoff_file))
    
    # 加载需水数据
    demand_file = forecast_files.get('demand_data', 'demand.json')
    demand_data = load_json(get_full_path(demand_file))
    
    # 加载目标约束数据
    obj_cons_file = forecast_files.get('obj_cons_data', 'obj_cons.json')
    obj_cons_data = load_json(get_full_path(obj_cons_file))
    
    logger.info("所有数据文件加载完成")
    
    # 返回所有原始数据
    return {
        'config': config,
        'runoff_data': runoff_data,
        'demand_data': demand_data,
        'obj_cons_data': obj_cons_data
    }

def load_water_supply_model(config_file: str = "config.json") -> Dict[str, Any]:
    """
    便捷函数：加载水库群供水调度优化模型的所有数据
    
    Args:
        config_file: 配置文件路径
        
    Returns:
        包含所有配置和原始数据的字典（未经处理的数据）
        如果需要处理后的数据，请使用 load_data.DataProcessor
    """
    # 设置项目根目录路径
    if not os.path.isabs(config_file):
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_file = os.path.join(project_root, config_file)
    
    # 使用initialize_environment加载配置
    config_files = {
        'config_filename': config_file,
        'filepath': os.path.dirname(config_file)
    }
    
    params = initialize_environment(config_files)
    
    # 加载所有原始JSON数据
    raw_data = load_all_json_data(config_file)
    
    # 合并参数（环境参数 + 原始数据）
    params.update(raw_data)
    
    return params

if __name__ == "__main__":
    try:
        script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_path = os.path.join(script_dir, "config.json")
        
        para = load_water_supply_model(config_path)
        logger.info("\n=== 数据加载完成 ===")
        
    except Exception as e:
        logger.error(f"加载失败: {e}")